package ru.geyes.screw.collections;

import ru.geyes.screw.closures.Function2V;
import ru.geyes.screw.closures.Function2P;
import ru.geyes.screw.closures.Function2B;
import ru.geyes.screw.closures.Function;

import java.util.HashMap;
import java.util.Collection;
import java.util.Map;

/**
 * @author alex
 *         date 17.05.2008
 *         time 22:17:16
 */
public class HashMapX<K, V> extends HashMap<K, V> implements MapX<K, V> {
    public HashMapX(Map<? extends K, ? extends V> m) {
        super(m);
    }

    public HashMapX(int initialCapacity, float loadFactor) {
        super(initialCapacity, loadFactor);
    }

    public HashMapX(int initialCapacity) {
        super(initialCapacity);
    }

    public HashMapX() {
        super();
    }

    public ListX<V> flatMap() {
        ListX<V> list = new ArrayListX<V>();
        for (V v : this.values()) {
            ListX<V> result;
            if (v instanceof Collection) {
                Collection c = (Collection)v;
                result = new ArrayListX<V>(c);
            } else if (v instanceof OptionX) {
                OptionX<V> c = (OptionX)v;
                if (c.isDefined())
                    result = ListXUtil.wrap(c.get());
                else
                    result = ListXUtil.emptyListX();
            } else {
                result = ListXUtil.wrap(v);
            }
            list.addAll(result);
        }
        return list;
    }

    public Option<V> getO(K key) {
        return Option.option(this.get(key));
    }

    public ListX<V> valuesX() {
        ListX<V> list = new ArrayListX<V>();
        list.addAll(this.values());
        return list;
    }

    public void forAll(Function2V<K, V> exec) {
        for (Map.Entry<K, V> entry : this.entrySet()) {
            K key = entry.getKey();
            V value = entry.getValue();
            exec.apply(key, value);
        }
    }

    public <S> MapX<K, S> map(Function2P<S, K, V> mapper) {
        MapX<K, S> newMap = MapXUtils.newHashMap();
        for (K k : this.keySet()) {
            newMap.put(k, mapper.apply(k, this.get(k)));
        }
        return newMap;
    }

    public MapX<K, V> filter(Function2B<K, V> filterF) {
        MapX<K, V> newMap = MapXUtils.newHashMap();
        for (K k : this.keySet()) {
            V v = this.get(k);
            if (filterF.apply(k, v)) newMap.put(k, v);
        }
        return newMap;
    }

    public V getOrPut(K key, Function<V> closure) {
        V value = this.get(key);
        if (value == null) {
            value = closure.apply();
            this.put(key, value);
        }
        return value;
    }
}
